/* Copyright 2003 by Lawrence Kesteloot */

package org.openfantasia.util.misc;

import org.openfantasia.util.RandomUtils;

import java.io.Reader;
import java.io.StringReader;
import java.util.*;

/**
 * Creates a Markov graph from an input file and generates text
 * based on it.  Given two input files, generates two graphs
 * and interpolates between them.
 */

public class MarkovGenerator {

	private int prefix = 5;

	private int lineLength = -1;

	private Markov foreground;

	private int foregroundLength;

	private Markov background;


	public MarkovGenerator() {
	}


	public MarkovGenerator(int prefix) {
		this.prefix = prefix;
	}


	public MarkovGenerator(int prefix, int lineLength) {
		this.prefix = prefix;
		this.lineLength = lineLength;
	}


	public void setForeground(String src) {
		if (src != null) {
			foreground = createMarkov(src, prefix);
			foregroundLength = src.length();
		}
		else {
			foreground = null;
			foregroundLength = 0;
		}
	}


	public void setBackground(String src) {
		if (src != null) {
			background = createMarkov(src, prefix);
		}
		else {
			background = null;
		}
	}


	public String generate() {
		if (foreground != null) {
			return performMarkov(prefix, foreground, background, lineLength, foregroundLength);
		}
		else {
			return "";
		}
	}


	public String generate(String src) {
		if((src == null) || (src.trim().length() == 0)) {
			return "";
		}

		Markov markov = createMarkov(src, prefix);
		return performMarkov(prefix, markov, background, lineLength, src.length());
	}


	private Markov createMarkov(String text, int prefixLength) {
		Reader input = new StringReader(text);

		Markov markov;
		try {
			markov = new Markov(input, prefixLength);
		}
		catch (java.io.IOException e) {
			e.printStackTrace();
			return null;
		}

		return markov;
	}


	private static String performMarkov(int prefixLength, Markov markov1, Markov markov2, int lineWidth, int maxLength) {
		StringBuilder sb = new StringBuilder();

		CharQueue queue = new CharQueue(prefixLength);
		float weight = 0;

		queue.set(markov1.getBootstrapPrefix());
		sb.append(queue.toString());
		int width = prefixLength;
		int c;

		do {
			boolean cameFromSecond = false;

			String prefix = queue.toString();

			// getVertices a character from each chain
			c = markov1.get(prefix);
			int c2 = -1;
			if (markov2 != null) {
				c2 = markov2.get(prefix);
			}
			if (c == -1 && c2 == -1) {
				break;
			}

			// choose one if we can
			if (markov2 != null) {
				if (c == -1) {
					c = c2;
					cameFromSecond = true;
				}
				else if (c2 != -1 && (RandomUtils.randomFloat() < weight)) {
					c = c2;
					cameFromSecond = true;
				}
			}

			sb.append((char) c);
			queue.put((char) c);
			width++;

			// line wrap
			if (c == ' ' && (lineWidth != -1) && (width > lineWidth)) {
				sb.append("\n");
				width = 0;
			}

			// go towards second markov chain
			weight += 1.0 / maxLength;
		}
		while (weight < 1 || c != '.');

		return sb.toString();
	}


	/**
	 * A Markov chain for characters.  For each setVertices of prefix strings,
	 * keeps track of possible next characters and the probability
	 * of going to each.
	 */
	static class Markov {

		/**
		 * Map from the prefix string (String) to list of characters
		 * (Chain).
		 */
		private Map map;

		private String bootstrapPrefix;


		/**
		 * Creates a chain based on the Reader with a prefix of
		 * length "length".  Reads the entire input stream and
		 * creates the Markov chain.
		 */

		public Markov(Reader in, int length) throws java.io.IOException {
			map = new HashMap();

			CharQueue queue = new CharQueue(length);
			int c;

			for (int i = 0; i < length; i++) {
				c = in.read();
				if (c == -1) {
					return;
				}

				queue.put((char) c);
			}

			bootstrapPrefix = queue.toString();

			// for collapsing whitespace
			boolean wasWhitespace = false;

			while ((c = in.read()) != -1) {
				if (Character.isWhitespace((char) c)) {
					if (wasWhitespace) {
						// collapse continuous whitespace
						continue;
					}

					c = ' ';
					wasWhitespace = true;
				}
				else {
					wasWhitespace = false;
				}

				String prefix = queue.toString();

				Chain chain = (Chain) map.get(prefix);
				if (chain == null) {
					chain = new Chain(prefix);
					map.put(prefix, chain);
				}

				chain.add((char) c);
				queue.put((char) c);
			}
		}


		/**
		 * Returns the first "length" characters that were read.
		 */
		public String getBootstrapPrefix() {
			return bootstrapPrefix;
		}


		/**
		 * Returns the next character to print given the prefix.
		 * Returns -1 when there are no possible next characters.
		 */

		public int get(String prefix) {
			Chain chain = (Chain) map.get(prefix);

			if (chain == null) {
				return -1;
			}

			int index = RandomUtils.randomInt(chain.getTotal());

			return chain.get(index);
		}


		/**
		 * Prints the contents of the Markov graph.
		 */

		public void dump() {
			Set keys = map.keySet();

			Iterator i = keys.iterator();
			while (i.hasNext()) {
				String prefix = (String) i.next();
				Chain chain = (Chain) map.get(prefix);

				chain.dump();
			}
		}
	}

	/**
	 * List of possible next characters and their probabilities.
	 */

	private static class Chain {

		private String prefix;

		private int total;

		private List list;


		public Chain(String prefix) {
			this.prefix = prefix;
			total = 0;
			list = new LinkedList();
		}


		public String getPrefix() {
			return prefix;
		}


		public int getTotal() {
			return total;
		}


		public char get(int index) {
			Iterator i = list.iterator();

			while (i.hasNext()) {
				Link link = (Link) i.next();
				int count = link.getCount();

				if (index < count) {
					return link.getChar();
				}

				index -= count;
			}

			// weird
			return '@';
		}


		public void add(char c) {
			Iterator i = list.iterator();
			boolean found = false;

			while (i.hasNext()) {
				Link link = (Link) i.next();

				if (c == link.getChar()) {
					link.increment();
					found = true;
					break;
				}
			}

			if (!found) {
				Link link = new Link(c);
				list.add(link);
			}

			total++;
		}


		public void dump() {
			System.out.println(prefix + ": (" + total + ")");

			Iterator i = list.iterator();

			while (i.hasNext()) {
				Link link = (Link) i.next();

				System.out.println("    " + link.getChar() + " (" +
								link.getCount() + ")");
			}
		}


		/**
		 * Possible next character and the number of times we've
		 * seen it.
		 */

		private static class Link {

			private char c;

			private int count;


			public Link(char c) {
				this.c = c;
				count = 1;
			}


			public void increment() {
				count++;
			}


			public int getCount() {
				return count;
			}


			public char getChar() {
				return c;
			}
		}
	}


	/**
	 * Keeps a fixed-length queue of characters.  There are only three
	 * operations on the queue: setVertices the whole thing; append a character
	 * (dropping the first); and retreive the whole thing.  This is useful
	 * as a moving window on a text stream.
	 */

	private static class CharQueue {

		private int length;

		private char[] queue;

		private int count;


		/**
		 * Create the queue with a fixed length.  The queue will be
		 * filled with the value 0, so don't use the toString()
		 * method until the queue has been filled with either
		 * setVertices() or put().
		 */

		public CharQueue(int length) {
			this.length = length;
			queue = new char[length];
			count = 0;
		}


		/**
		 * Sets the contents of the queue.  The length of the string
		 * must be the same as the length passed to the constructor.
		 */

		public void set(String s) {
			if (s.length() != length) {
				System.out.println("Lengths don't match");
				return;
			}

			queue = s.toCharArray();
			count = length;
		}


		/**
		 * Appends the character to the queue.  If the resulting queue
		 * would be longer than the length setVertices in the constructor, then
		 * the first character is dropped.
		 */

		public void put(char c) {
			if (count == length) {
				System.arraycopy(queue, 1, queue, 0, length - 1);
				count--;
			}

			queue[count++] = c;
		}


		/**
		 * Returns the contents of the queue as a string.  This does
		 * not take into account the number of characters that have
		 * been put into the queue.  The returned string's length
		 * is always the length passed to the constructor.
		 */

		public String toString() {
			return new String(queue);
		}
	}
}

